{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"collapsed_sections":["UeFSNswvparA"]},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"}},"cells":[{"cell_type":"markdown","source":["Copyright 2022 Google LLC. SPDX-License-Identifier: Apache-2.0\n","\n","Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n","\n","https://www.apache.org/licenses/LICENSE-2.0\n","\n","Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."],"metadata":{"id":"D4K98vDLEfi8"}},{"cell_type":"markdown","source":["# Language Model Program Examples\n","\n","This notebook is a part of the open-source code release associated with the paper:\n","\n","[Code as Policies: Language Model Programs for Embodied Control](https://code-as-policies.github.io/)\n","\n","This notebook accompanies examples given in the Method section of the paper.\n","\n","1) Please obtain an OpenAI API Key here:\n","https://openai.com/blog/openai-api/\n","\n","2) Gain Codex access by joining the waitlist here:\n","https://openai.com/blog/openai-codex/\n","\n","Once you have Codex access you can use `code-davinci-002` as the `model_name`. Using the GPT-3 model (`text-dainvci-002`) is also ok, but performance won't be as good (there will be more code logic errors)."],"metadata":{"id":"3AbXc_pfEjSl"}},{"cell_type":"code","source":["openai_api_key = 'YOUR KEY HERE'\n","model_name = 'code-davinci-002' # 'text-davinci-002'"],"metadata":{"id":"V-jOpomPG4fX"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Setup"],"metadata":{"id":"f_bHZvLZoNxH"}},{"cell_type":"code","execution_count":null,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"0JPbDmNeoIfk","executionInfo":{"status":"ok","timestamp":1662589626947,"user_tz":240,"elapsed":15161,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"59f51cce-49b0-4c9e-ba41-31475c126d0e"},"outputs":[{"output_type":"stream","name":"stdout","text":["Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n","Collecting openai\n","  Downloading openai-0.23.0.tar.gz (43 kB)\n","\u001b[K     |████████████████████████████████| 43 kB 1.6 MB/s \n","\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n","  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n","    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n","Requirement already satisfied: openpyxl>=3.0.7 in /usr/local/lib/python3.7/dist-packages (from openai) (3.0.10)\n","Requirement already satisfied: requests>=2.20 in /usr/local/lib/python3.7/dist-packages (from openai) (2.23.0)\n","Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from openai) (4.64.0)\n","Requirement already satisfied: pandas>=1.2.3 in /usr/local/lib/python3.7/dist-packages (from openai) (1.3.5)\n","Collecting pandas-stubs>=1.1.0.11\n","  Downloading pandas_stubs-1.2.0.62-py3-none-any.whl (163 kB)\n","\u001b[K     |████████████████████████████████| 163 kB 9.3 MB/s \n","\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from openai) (1.21.6)\n","Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from openai) (4.1.1)\n","Requirement already satisfied: et-xmlfile in /usr/local/lib/python3.7/dist-packages (from openpyxl>=3.0.7->openai) (1.1.0)\n","Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=1.2.3->openai) (2.8.2)\n","Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas>=1.2.3->openai) (2022.2.1)\n","Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.7.3->pandas>=1.2.3->openai) (1.15.0)\n","Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (2.10)\n","Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (2022.6.15)\n","Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (1.24.3)\n","Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests>=2.20->openai) (3.0.4)\n","Building wheels for collected packages: openai\n","  Building wheel for openai (PEP 517) ... \u001b[?25l\u001b[?25hdone\n","  Created wheel for openai: filename=openai-0.23.0-py3-none-any.whl size=54478 sha256=c51285f42e11af3e1778187062418b117743adb381602e9e24f6536e6c0c37e2\n","  Stored in directory: /root/.cache/pip/wheels/70/d5/31/f9f67660319d89e4f54501d27b1e90f88a3309c42ea4fd734c\n","Successfully built openai\n","Installing collected packages: pandas-stubs, openai\n","Successfully installed openai-0.23.0 pandas-stubs-1.2.0.62\n"]}],"source":["! pip install openai\n","\n","import openai\n","openai.api_key = openai_api_key"]},{"cell_type":"code","source":["def lmp(base_prompt, query, stop_tokens=None, query_kwargs=None):\n","    new_prompt = f'{base_prompt}\\n{query}'\n","\n","    use_query_kwargs = {\n","        'engine': model_name,\n","        'max_tokens': 512,\n","        'temperature': 0,\n","    }\n","    if query_kwargs is not None:\n","      use_query_kwargs.update(query_kwargs)\n","\n","    response = openai.Completion.create(\n","        prompt=new_prompt, stop=stop_tokens, **use_query_kwargs\n","    )['choices'][0]['text'].strip()\n","\n","    print(query)\n","    print(response)\n","\n","    return response"],"metadata":{"id":"jq8rs8c5oPMc"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Simple LMPs"],"metadata":{"id":"hLiqQ2VUobAY"}},{"cell_type":"markdown","source":["## Pure Python"],"metadata":{"id":"PqpoB26Ioeww"}},{"cell_type":"code","source":["prompt_pure_python = '''\n","# Python script\n","# get the variable a.\n","ret_val = a\n","'''.strip()"],"metadata":{"id":"gCTH9KPtogbw"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["_ = lmp(prompt_pure_python, '# find the sum of variables a and b.', ['#'])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"WynEiOt_Rck4","executionInfo":{"status":"ok","timestamp":1662488939333,"user_tz":240,"elapsed":919,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"6b561a70-3404-48f6-d494-bfcde190748b"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# find the sum of variables a and b.\n","ret_val = a + b\n"]}]},{"cell_type":"code","source":["_ = lmp(prompt_pure_python, '# find the sum of numbers in a list called values.', ['#'])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"AR-dNTgRp0DO","executionInfo":{"status":"ok","timestamp":1660759746864,"user_tz":240,"elapsed":1373,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"02039b9b-adda-4e64-a77d-8e2cf45ebad8"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# find the sum of numbers in a list called values.\n","ret_val = sum(values)\n"]}]},{"cell_type":"code","source":["_ = lmp(prompt_pure_python, '# find the difference between the max and min numbers in a list called xs.', ['#'])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"5pU6wU3cqlbO","executionInfo":{"status":"ok","timestamp":1660759749902,"user_tz":240,"elapsed":1685,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"0e1cf553-91c7-4d83-92fe-2c686f69f7a5"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# find the difference between the max and min numbers in a list called xs.\n","ret_val = max(xs) - min(xs)\n"]}]},{"cell_type":"code","source":["_ = lmp(prompt_pure_python, '# see if any number is divisible by 3 in a list called xs.', ['#'])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"A8jW86PJqwX5","executionInfo":{"status":"ok","timestamp":1660771878438,"user_tz":240,"elapsed":1480,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"3f6e8861-9e2b-4f0f-d9f0-03b4f34ececc"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# see if any number is divisible by 3 in a list called xs.\n","ret_val = any(x % 3 == 0 for x in xs)\n"]}]},{"cell_type":"markdown","source":["## Using Context"],"metadata":{"id":"w3uCvfdyogiq"}},{"cell_type":"code","source":["prompt_context = '''\n","objects = ['green block', 'green bowl', 'yellow block', 'yellow bowl']\n","# the yellow block.\n","ret_val = 'yellow block'\n","# the blocks.\n","ret_val = ['green block', 'yellow block']\n","'''.strip()"],"metadata":{"id":"2SUdOnJBonGU"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["context = \"objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\"\n","query = '# the bowls.'\n","\n","print(context)\n","_ = lmp(f'{prompt_context}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"m90wMsv0blpA","executionInfo":{"status":"ok","timestamp":1660772082794,"user_tz":240,"elapsed":928,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"8fedf59a-390c-4d33-eec0-11a565c03e05"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\n","# the bowls.\n","ret_val = ['blue bowl', 'red bowl']\n"]}]},{"cell_type":"code","source":["context = \"objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\"\n","query = '# sea-colored block.'\n","\n","print(context)\n","_ = lmp(f'{prompt_context}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"Ib4qYMXPudXL","executionInfo":{"status":"ok","timestamp":1660772086408,"user_tz":240,"elapsed":950,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"b6cb73e3-04b5-4863-9805-d2f66f340a9a"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\n","# sea-colored block.\n","ret_val = 'blue block'\n"]}]},{"cell_type":"code","source":["context = '''\n","objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\n","# sea-colored block.\n","ret_val = 'blue block'\n","objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\n","'''.strip()\n","query = '# the other block.'\n","\n","print(context)\n","_ = lmp(f'{prompt_context}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"-3tRUPdKwvZz","executionInfo":{"status":"ok","timestamp":1660839794184,"user_tz":240,"elapsed":1817,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"bb40fa37-6092-4a5e-a3a3-4ea4b3cda723"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\n","# sea-colored block.\n","ret_val = 'blue block'\n","objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\n","# the other block.\n","ret_val = 'red block'\n"]}]},{"cell_type":"markdown","source":["## Using 3rd Party Libraries"],"metadata":{"id":"zNpyKf9honRJ"}},{"cell_type":"code","source":["prompt_3lib = '''\n","import numpy as np\n","# move all points in pts_np toward the right.\n","ret_val = pts_np + [0.3, 0]\n","# move a pt_np toward the top.\n","ret_val = pt_np + [0, 0.3]\n","'''.strip()"],"metadata":{"id":"MYVbZ9bIopHm"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["query = '# get the left most point in pts_np.'\n","\n","_ = lmp(prompt_3lib, query, ['#'])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"TfKxQ70MSaw8","executionInfo":{"status":"ok","timestamp":1660832123198,"user_tz":240,"elapsed":2733,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"42883e86-e5be-4e00-8c00-c3ab3feb59dc"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# get the left most point in pts_np.\n","left_most = pts_np[np.argmin(pts_np[:, 0])]\n"]}]},{"cell_type":"code","source":["query = '# get the center of pts_np.'\n","\n","_ = lmp(prompt_3lib, query, ['#'])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"_LVXn0VCdcig","executionInfo":{"status":"ok","timestamp":1660832040281,"user_tz":240,"elapsed":2929,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"5d7a0ac1-d4c5-4785-ad64-da9c8638018d"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# get the center of pts_np.\n","center = np.mean(pts_np, axis=0)\n"]}]},{"cell_type":"code","source":["query = '# the closest point in pts_np to pt_np.'\n","\n","_ = lmp(prompt_3lib, query, ['#'])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"wDtLNi5pdvAH","executionInfo":{"status":"ok","timestamp":1660832129955,"user_tz":240,"elapsed":3155,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"2e4c0b1a-e4c8-4722-d7db-abc8eb238b78"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# the closest point in pts_np to pt_np.\n","ret_val = pts_np[np.argmin(np.linalg.norm(pts_np - pt_np, axis=1))]\n"]}]},{"cell_type":"markdown","source":["## Using 1st Party Libraries"],"metadata":{"id":"g5vpw9m_opST"}},{"cell_type":"code","source":["prompt_1lib = '''\n","from utils import get_pos, put_first_on_second\n","objects = ['gray block', 'gray bowl']\n","# put the gray block on the gray bowl.\n","put_first_on_second('gray block', 'gray bowl')\n","objects = ['purple block', 'purple bowl']\n","# move the purple bowl toward the left.\n","target_pos = get_pos('purple bowl') + [-0.3, 0]\n","put_first_on_second('purple bowl', target_pos)\n","'''.strip()"],"metadata":{"id":"WWwkEvYzork6"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["context = \"objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\"\n","query = '# move the red block a bit to the right.'\n","\n","print(context)\n","_ = lmp(f'{prompt_1lib}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"RONmfY-ETcMf","executionInfo":{"status":"ok","timestamp":1660773658350,"user_tz":240,"elapsed":2546,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"44fc2d84-a191-4f33-bd5e-3b74ca368755"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\n","# move the red block a bit to the right.\n","target_pos = get_obj_pos('red block') + [0.1, 0]\n","put_first_on_second('red block', target_pos)\n"]}]},{"cell_type":"code","source":["context = \"objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\"\n","query = '# put the blue block on the bowl with the same color.'\n","\n","print(context)\n","_ = lmp(f'{prompt_1lib}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"DWVCGpPoy7WC","executionInfo":{"status":"ok","timestamp":1660773773456,"user_tz":240,"elapsed":1794,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"d8b34bf3-54fd-4dc5-9967-588a4cd49881"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['blue bowl', 'red block', 'red bowl', 'blue block']\n","# put the blue block on the bowl with the same color.\n","put_first_on_second('blue block', 'blue bowl')\n"]}]},{"cell_type":"markdown","source":["## Combined Examples"],"metadata":{"id":"lmyR8qsoT6qy"}},{"cell_type":"code","source":["prompt_combined = '''\n","import numpy as np\n","from utils import get_pos, put_first_on_second\n","objects = ['cyan block', 'cyan bowl', 'pink bowl']\n","# put the cyan block in cyan bowl.\n","put_first_on_second('cyan block', 'cyan bowl')\n","objects = ['gray block', 'silver block', 'gray bowl']\n","# place the top most block on the gray bowl.\n","names = ['gray block', 'silver block']\n","positions = np.array([get_pos(name) for name in names])\n","name = names[np.argmax(positions[:,1])]\n","put_first_on_second(name, 'gray bowl')\n","objects = ['purple block', 'purple bowl']\n","# put the purple bowl to the left of the purple block.\n","target_pos = get_pos('purple block') + [-0.3, 0]\n","put_first_on_second('purple bowl', target_pos)\n","'''.strip()"],"metadata":{"id":"sahRYGJCWzYJ"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["context = \"objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\"\n","query = '# move the left most bowl toward the right.'\n","\n","print(context)\n","_ = lmp(f'{prompt_combined}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"HQyGNJM2UCkQ","executionInfo":{"status":"ok","timestamp":1660833464404,"user_tz":240,"elapsed":7012,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"120c6770-f191-41fb-da28-58b8129671e8"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\n","# move the left most bowl toward the right.\n","names = ['red bowl', 'blue bowl']\n","positions = np.array([get_pos(name) for name in names])\n","name = names[np.argmin(positions[:,0])]\n","target_pos = get_pos(name) + [0.3, 0]\n","put_first_on_second(name, target_pos)\n"]}]},{"cell_type":"code","source":["context = \"objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\"\n","query = '# place the blocks in bowls with their colors.'\n","\n","print(context)\n","_ = lmp(f'{prompt_combined}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"0zUt121Cgw0z","executionInfo":{"status":"ok","timestamp":1660833578803,"user_tz":240,"elapsed":3553,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"75e3e170-d006-47cf-bbbd-4126735bd27e"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\n","# place the blocks in bowls with their colors.\n","put_first_on_second('red block', 'red bowl')\n","put_first_on_second('blue block', 'blue bowl')\n"]}]},{"cell_type":"code","source":["context = \"objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\"\n","query = '# move the red block to the middle of the bowls.'\n","\n","print(context)\n","_ = lmp(f'{prompt_combined}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"OM4PcAVuijeQ","executionInfo":{"status":"ok","timestamp":1660833763311,"user_tz":240,"elapsed":3663,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"8739f01c-274f-4c3a-b6a3-e0f3d89223ab"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\n","# move the red block to the middle of the bowls.\n","target_pos = np.mean(np.array([get_pos('blue bowl'), get_pos('red bowl')]), axis=0)\n","put_first_on_second('red block', target_pos)\n"]}]},{"cell_type":"markdown","source":["# Complex LMPs"],"metadata":{"id":"Lfvev56boc3J"}},{"cell_type":"markdown","source":["## Control Flows"],"metadata":{"id":"wLg4-TEXpZL5"}},{"cell_type":"code","source":["prompt_ctrl = '''\n","from utils import get_pos, put_first_on_second\n","# move the orange block toward the top.\n","target_pos = get_pos('orange block') + [0, 0.3]\n","put_first_on_second('orange block', target_pos)\n","# move the purple bowl toward the left.\n","target_pos = get_pos('purple bowl') + [-0.3, 0]\n","put_first_on_second('purple bowl', target_pos)\n","'''"],"metadata":{"id":"S5HEf5gIod8h"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["context = \"objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\"\n","query = '# while the red block is to the left of the blue bowl, move it to the right 5cm at a time.'\n","\n","_ = lmp(f'{prompt_ctrl}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"oqFojGZ1dCYv","executionInfo":{"status":"ok","timestamp":1660835466035,"user_tz":240,"elapsed":6242,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"55b36b11-0987-4101-9177-6dc5a9451e2f"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# while the red block is to the left of the blue bowl, move it to the right 5cm at a time.\n","while get_pos('red block')[0] < get_pos('blue bowl')[0]:\n","    target_pos = get_pos('red block') + [0.05, 0]\n","    put_first_on_second('red block', target_pos)\n"]}]},{"cell_type":"code","source":["context = \"objects = ['yellow block', 'green bowl', 'green block', 'yellow bowl']\"\n","query = \"# move the yellow block toward its bowl 1cm at a time until their distance is less than 5cm apart.\"\n","\n","_ = lmp(f'{prompt_ctrl}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"zVFAg_P_uRdY","executionInfo":{"status":"ok","timestamp":1660836378958,"user_tz":240,"elapsed":6227,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"cc76e284-822b-445f-b8d1-16e2d93fcdb1"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# move the yellow block toward its bowl 1cm at a time until their distance is less than 5cm apart.\n","target_pos = get_pos('yellow bowl')\n","while np.linalg.norm(get_pos('yellow block') - target_pos) > 0.05:\n","    put_first_on_second('yellow block', target_pos)\n","    target_pos += [0, 0.01]\n"]}]},{"cell_type":"markdown","source":["## Calling Other LMPs"],"metadata":{"id":"UeFSNswvparA"}},{"cell_type":"code","source":["prompt_modular_ui = '''\n","import numpy as np\n","from utils import get_pos, put_first_on_second, parse_obj\n","objects = ['yellow block', 'yellow bowl', 'gray block', 'gray bowl']\n","# move the sun colored block toward the left.\n","block_name = parse_obj('sun colored block')\n","target_pos = get_pos(block_name) + [-0.3, 0]\n","put_first_on_second(block_name, target_pos)\n","objects = ['white block', 'white bowl', 'yellow block', 'yellow bowl']\n","# place the block closest to the blue bowl on the other bowl.\n","block_name = parse_obj('the block closest to the blue bowl', f'objects = {objects}')\n","bowl_name = parse_obj('a bowl other than the blue bowl', f'objects = {objects}')\n","put_first_on_second(block_name, bowl_name)\n","'''.strip()\n","\n","prompt_modular_parse_obj = '''\n","import numpy as np\n","from utils import get_pos\n","objects = ['brown bowl', 'green block', 'brown block', 'green bowl']\n","# the blocks.\n","ret_val = ['brown block', 'green block']\n","# the sky colored block.\n","ret_val = 'blue block'\n","objects = ['orange block', 'cyan block', 'purple bowl', 'gray bowl']\n","# the right most block.\n","block_names = ['orange block', 'cyan block']\n","block_positions = np.array([get_pos(block_name) for block_name in block_names])\n","right_block_name = block_names[np.argmax(block_positions[:, 0])]\n","ret_val = right_block_name\n","'''.strip()"],"metadata":{"id":"upuOWzkupe9W"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["context = \"objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\"\n","query = '# while the left most block is the red block, move it toward the right.'\n","\n","print(context)\n","_ = lmp(f'{prompt_modular_ui}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"1e3AUta4d8Xn","executionInfo":{"status":"ok","timestamp":1660764659496,"user_tz":240,"elapsed":5425,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"ad12ad9a-e2b9-4b50-f5df-7af88a539439"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\n","# while the left most block is the red block, move it toward the right.\n","block_name = parse_obj('the left most block')\n","while block_name == 'red block':\n","    target_pos = get_pos(block_name) + [0.3, 0]\n","    put_first_on_second(block_name, target_pos)\n","    block_name = parse_obj('the left most block')\n"]}]},{"cell_type":"code","source":["context = \"objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\"\n","query = '# the left most block.'\n","\n","print(context)\n","_ = lmp(f'{prompt_modular_parse_obj}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"lTOB3IUweEEs","executionInfo":{"status":"ok","timestamp":1660764860471,"user_tz":240,"elapsed":4317,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"7556dad3-65fc-453f-dc34-2e1b2a48656f"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\n","# the left most block.\n","block_names = ['red block', 'blue block']\n","block_positions = np.array([get_pos(block_name) for block_name in block_names])\n","left_block_name = block_names[np.argmin(block_positions[:, 0])]\n","ret_val = left_block_name\n"]}]},{"cell_type":"markdown","source":["## Function-Generating LMPs"],"metadata":{"id":"c0ppNd3zpfXv"}},{"cell_type":"code","source":["prompt_f_gen = '''\n","import numpy as np\n","from utils import get_pos, get_obj_bbox_xyxy\n","# define function: total = get_total(xs=numbers).\n","def get_total(xs):\n","    return np.sum(xs)\n","'''.strip()"],"metadata":{"id":"dXAbRNAHpikl"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["query = '# define function: get_objs_bigger_than_area_th(obj_names, bbox_area_th).'\n","\n","_ = lmp(f'{prompt_f_gen}', query, ['# define', '# example'])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"ZnFo0QQRkFa4","executionInfo":{"status":"ok","timestamp":1660836781661,"user_tz":240,"elapsed":6767,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"7f0708de-8984-4212-f253-b63d65bb56db"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# define function: get_objs_bigger_than_area_th(obj_names, bbox_area_th).\n","def get_objs_bigger_than_area_th(obj_names, bbox_area_th):\n","    return [obj_name for obj_name in obj_names if get_obj_bbox_area(obj_name) > bbox_area_th]\n"]}]},{"cell_type":"markdown","source":["## Hierarchical Function-Generating LMPs"],"metadata":{"id":"lU9WRbh2GOpS"}},{"cell_type":"code","source":["query = '# define function: get_obj_bbox_area(obj_name).'\n","\n","_ = lmp(prompt_f_gen, query, ['# define', '# example'])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"-4aWySGjl7Rd","executionInfo":{"status":"ok","timestamp":1660836894097,"user_tz":240,"elapsed":6845,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"5b1ff0c2-d36e-490c-d343-81cb33c78462"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# define function: get_obj_bbox_area(obj_name).\n","def get_obj_bbox_area(obj_name):\n","    x1, y1, x2, y2 = get_obj_bbox_xyxy(obj_name)\n","    return (x2 - x1) * (y2 - y1)\n"]}]},{"cell_type":"code","source":["prompt_hier_f_gen = '''\n","import numpy as np\n","from utils import get_obj_pos, get_obj_bbox_xyxy\n","# define function: total = get_total(xs=numbers).\n","def get_total(xs):\n","    return np.sum(xs)\n","# define function: pt_np = move_pt_left(pt_np, dist).\n","def move_pt_left(pt_np, dist):\n","    delta = np.array([-dist, 0])\n","    return translate_pt_np(pt_np, delta=delta)\n","'''.strip()"],"metadata":{"id":"-mtqE9qbGQl_"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["query = '# define function: get_obj_bbox_area(obj_name).'\n","\n","_ = lmp(prompt_hier_f_gen, query, ['# define', '# example'])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"xqAoMBuUkcH2","executionInfo":{"status":"ok","timestamp":1660836512561,"user_tz":240,"elapsed":4227,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"9d5b2c7f-ff14-4e22-91c6-bfdc0ea6c7ba"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# define function: get_obj_bbox_area(obj_name).\n","def get_obj_bbox_area(obj_name):\n","    bbox_xyxy = get_obj_bbox_xyxy(obj_name)\n","    return get_bbox_area(bbox_xyxy)\n"]}]},{"cell_type":"code","source":["query = '# define function: area = get_bbox_area(bbox_xyxy).'\n","\n","_ = lmp(prompt_hier_f_gen, query, ['# define', '# example'])"],"metadata":{"id":"mTJyh1LuGo7g","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1660753295019,"user_tz":240,"elapsed":3425,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"4f4f87d7-584d-4f1f-e000-54833cf19979"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["# define function: area = get_bbox_area(bbox_xyxy).\n","def get_bbox_area(bbox_xyxy):\n","    x1, y1, x2, y2 = bbox_xyxy\n","    return (x2 - x1) * (y2 - y1)\n"]}]},{"cell_type":"markdown","source":["## Combined Example"],"metadata":{"id":"H70qFhsCgS37"}},{"cell_type":"code","source":["prompt_combined_parse_pos = '''\n","from utils import get_obj_pos\n","\n","# right of the red block.\n","ret_val = get_obj_pos('red block') + np.array([0.3, 0])\n","# a bit below the cyan bowl.\n","ret_val = get_obj_pos('cyan bowl') + np.array([0, -0.1])\n","# a point between the cyan block and purple bowl.\n","pos = get_mid_point_np(pt0=get_obj_pos('cyan block'), pt1=get_obj_pos('purple bowl'))\n","ret_val = pos\n","# the corner closest to the yellow block.\n","corner_positions = get_corner_positions()\n","closest_corner_idx = get_closest_idx(points=corner_positions, point=get_obj_pos('yellow block'))\n","closest_corner_pos = corner_positions(closest_corner_idx)\n","ret_val = closest_corner_pos\n","'''.strip()\n","\n","prompt_combined_parse_obj = '''\n","import numpy as np\n","from utils import get_obj_pos, get_objs_bigger_than_area_th\n","\n","objects = ['brown bowl', 'green block', 'brown block', 'green bowl', 'blue bowl', 'blue block']\n","# the blocks.\n","ret_val = ['brown block', 'blue block']\n","# the sky colored block.\n","ret_val = 'blue block'\n","objects = ['blue block', 'cyan block', 'purple bowl', 'gray bowl', 'brown bowl', 'pink block']\n","# the right most block.\n","block_names = ['cyan block', 'pink block', 'blue block']\n","block_positions = np.array([get_obj_pos(block_name) for block_name in block_names])\n","right_block_name = block_names[np.argmax(block_positions[:, 0])]\n","ret_val = right_block_name\n","objects = ['blue block', 'cyan block', 'purple bowl', 'brown bowl', 'purple block']\n","# blocks above the brown bowl.\n","block_names = ['blue block', 'cyan block', 'purple block']\n","brown_bowl_pos = get_obj_pos('brown bowl')\n","use_block_names = [name for name in block_names if get_obj_pos(name)[1] > brown_bowl_pos[1]]\n","ret_val = use_block_names\n","'''.strip()\n","\n","prompt_combined_ui = '''\n","import numpy as np\n","from utils import get_obj_pos, put_first_on_second, parse_obj\n","\n","objects = ['green block', 'green bowl', 'yellow block', 'yellow bowl']\n","# move the bowls below the yellow block upwards.\n","bowl_names = parse_obj('the bowls below the yellow block', f'objects = {objects}')\n","for bowl_name in bowl_names:\n","  target_pos = get_obj_pos(bowl_name) + np.array([0, 0.1])\n","  put_first_on_second(bowl_name, target_pos)\n","objects = ['white block', 'white bowl', 'yellow block', 'yellow bowl']\n","# if the banana colored bowl is the top most bowl, then place it on the white bowl.\n","bowl_name = parse_obj('the banana colored bowl', f'objects = {objects}')\n","top_bowl_name = parse_obj('the top most bowl', f'objects = {objects}')\n","if bowl_name == top_bowl_name:\n","  put_first_on_second(bowl_name, 'white bowl')\n","objects = ['blue block', 'blue bowl', 'red block', 'red bowl']\n","# place the block closest to the blue bowl on the other bowl.\n","block_name = parse_obj('the block closest to the blue bowl', f'objects = {objects}')\n","bowl_name = parse_obj('a bowl other than the blue bowl', f'objects = {objects}')\n","put_first_on_second(block_name, 'red bowl')\n","'''.strip()"],"metadata":{"id":"ZbC8ABgOgTi3"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["context = \"objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\"\n","query = '# while there are blocks with area bigger than 0.2 that are left of the red bowl, move them toward the right.'\n","\n","print(context)\n","_ = lmp(f'{prompt_combined_ui}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"DoxawWH2ndBB","executionInfo":{"status":"ok","timestamp":1660838726454,"user_tz":240,"elapsed":7025,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"fef1576a-a084-449c-b496-f13ec15ae873"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\n","# while there are blocks with area bigger than 0.2 that are left of the red bowl, move them toward the right.\n","block_names = parse_obj('blocks with area bigger than 0.2 that are left of the red bowl', f'objects = {objects}')\n","while len(block_names) > 0:\n","  for block_name in block_names:\n","    target_pos = get_obj_pos(block_name) + np.array([0.1, 0])\n","    put_first_on_second(block_name, target_pos)\n","  block_names = parse_obj('blocks with area bigger than 0.2 that are left of the red bowl', f'objects = {objects}')\n"]}]},{"cell_type":"code","source":["context = \"objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\"\n","query = '# blocks with area bigger than 0.2 that are left of the red bowl.'\n","\n","print(context)\n","_ = lmp(f'{prompt_combined_parse_obj}\\n{context}', query, ['#', 'objects = ['])"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"gjs2ICvYpeGK","executionInfo":{"status":"ok","timestamp":1660838886623,"user_tz":240,"elapsed":8633,"user":{"displayName":"Jacky Liang","userId":"05524594537381871813"}},"outputId":"b4c7d735-ab53-4767-942e-1a2040264a42"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["objects = ['red block', 'blue bowl', 'blue block', 'red bowl']\n","# blocks with area bigger than 0.2 that are left of the red bowl.\n","block_names = ['red block', 'blue block']\n","red_bowl_pos = get_obj_pos('red bowl')\n","use_block_names = [name for name in block_names if get_obj_pos(name)[0] < red_bowl_pos[0]]\n","use_block_names = get_objs_bigger_than_area_th(use_block_names, 0.2)\n","ret_val = use_block_names\n"]}]}]}